import random
import string
import logging
import paramiko
from celery import shared_task
from django.utils import timezone
from .models import Host, HostStat
from django.db.models import Count
from django_cryptography.fields import encrypt

# 初始化日志
logger = logging.getLogger(__name__)


@shared_task
def change_root_passwords():
    hosts = Host.objects.all()
    ssh_username = 'root'  # 假设是root用户登录
    ssh_port = 22  # SSH 默认端口
    for host in hosts:
        try:
            new_password = ''.join(
                random.choices(
                    string.ascii_letters + string.digits + '!@#$%^&*',
                    k=12
                )
            )

            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            ssh.connect(
                hostname=host.hostname,
                port=ssh_port,
                username=ssh_username,
                password=host.root_password.decrypt()
            )

            stdin, stdout, stderr = ssh.exec_command(f'echo "root:{new_password}" | chpasswd')

            output = stdout.read().decode()
            error = stderr.read().decode()

            if error:
                logger.error(f"Error changing password on {host.hostname}: {error}")
            else:
                logger.info(f"Password changed successfully on {host.hostname}")
                host.root_password = encrypt(new_password)
                host.last_password_change = timezone.now()
                host.save()

            ssh.close()
        except Exception as e:
            logger.error(f"Failed to connect or change password on {host.hostname}: {e}")




@shared_task
def generate_host_stats():
    """每天00:00生成主机统计"""
    today = timezone.now().date()

    stats = Host.objects.values('idc__city', 'idc').annotate(
        count=Count('id')
    )

    for stat in stats:
        HostStat.objects.create(
            city_id=stat['idc__city'],
            idc_id=stat['idc'],
            host_count=stat['count'],
            stat_date=today
        )

    return f"生成 {len(stats)} 条统计记录"
